package rbtree;

public class RBTree {

    public RBTreeNode root;// 根节点

    static class RBTreeNode{
        public RBTreeNode left;
        public RBTreeNode right;
        public RBTreeNode parent;
        public int val;
        public Color color;

        public RBTreeNode(int val) {
            this.val = val;
            // 新创建的节点, 颜色默认为红色
            this.color = Color.RED;
        }
    }


    // 插入元素
    public boolean insert(int val){
        // 先按照二叉搜索树的方式插入元素
        RBTreeNode node = new RBTreeNode(val);
        if (root == null) {
            root = node;
            root.color = Color.BLACK;
            return true;
        }
        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while (cur != null){
            if (cur.val > val) {
                // 在左子树中继续找
                parent = cur;
                cur = cur.left;
            }else if (cur.val < val){
                // 在右子树中继续找
                parent = cur;
                cur = cur.right;
            }else {
                return false;// 找到节点,插不进去了
            }
        }
        // 找到对应位置, 插入节点
        if (parent.val > val){
            parent.left = node;
        }else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;

        // 调整颜色
        while (parent != null && parent.color == Color.RED){
            RBTreeNode grandFather = parent.parent;
            RBTreeNode uncle = null;

            if (parent == grandFather.left){
                uncle = grandFather.right;
                if (uncle != null && uncle.color == Color.RED){
                    parent.color = Color.BLACK;
                    uncle.color = Color.BLACK;
                    grandFather.color = Color.RED;

                    // 继续向上修改
                    cur = grandFather;
                    parent = cur.parent;
                }else {
                    // uncle不存在, 或者为黑色

                    // 情况三：cur在 parent的右侧  --> 左旋变成情况二
                    if (cur == parent.right){
                        rotateLeft(parent, cur);
                        RBTreeNode temp = parent;
                        parent = cur;
                        cur = temp;
                    }

                    // 情况二：cur在 parent的左边
                    rotateRight(grandFather, parent);
                    grandFather.color = Color.RED;
                    parent.color = Color.BLACK;
                }
            }else {
                // parent在 grandFather的右边
                uncle = grandFather.left;
                if (uncle != null && uncle.color == Color.RED){
                    parent.color = Color.BLACK;
                    uncle.color = Color.BLACK;
                    grandFather.color = Color.RED;

                    // 继续向上修改
                    cur = grandFather;
                    parent = cur.parent;
                }else {
                    // uncle不存在, 或者为黑色

                    // 情况三：cur在 parent的右侧  --> 左旋变成情况二
                    if (cur == parent.left){
                        rotateRight(parent, cur);
                        RBTreeNode temp = parent;
                        parent = cur;
                        cur = temp;
                    }

                    // 情况二：cur在 parent的左边
                    rotateLeft(grandFather, parent);
                    grandFather.color = Color.RED;
                    parent.color = Color.BLACK;
                }
            }
        }

        root.color = Color.BLACK;
        return true;
    }

    // 右单旋
    public void rotateRight(RBTreeNode parent, RBTreeNode cur){
        RBTreeNode subLeft = cur;
        RBTreeNode subRight = cur.right;
        RBTreeNode pParent = parent.parent;

        subLeft.right = parent;
        parent.left = subRight;

        parent.parent = subLeft;
        if (subRight != null){
            subRight.parent = parent;
        }

        if (pParent == null){
            // parent为根节点
            root = subLeft;
            subLeft.parent = null;
        }else {
            if (parent == pParent.left){
                pParent.left = subLeft;
            }else {
                pParent.right = subLeft;
            }
            subLeft.parent = pParent;
        }
    }

    // 左单旋
    public void rotateLeft(RBTreeNode parent, RBTreeNode cur){
        RBTreeNode subLeft = cur.left;
        RBTreeNode subRight = cur;
        RBTreeNode pParent = parent.parent;

        parent.right = subLeft;
        cur.left = parent;
        parent.parent = subRight;
        if (subLeft != null){
            subLeft.parent = parent;
        }

        if (pParent == null){
            root = subRight;
            root.parent = null;
        }else {
            if (pParent.left == parent){
                pParent.left = subRight;
            }else {
                pParent.right = subRight;
            }
            subRight.parent = pParent;
        }
    }

    /**
     * 判断当前树是否为红黑树
     */
    public boolean isRbTree(){

        if (root == null) return true;
        // 根节点是黑的
        if (root.color != Color.BLACK) {
            System.out.println("根节点不是黑色的!");
            return false;
        }
        // 要判断是否有两个连续节点， 先遍历到一个红色节点，看他的父节点是否为红色
        // 判断 每条路径是否黑色节点一样多
        int blackNum = 0;
        RBTreeNode node = root;
        while (node != null){
            if (node.color == Color.BLACK){
                blackNum++;
            }
            node = node.left;
        }

        return checkColor(root) && checkBlackNumber(root, 0, blackNum);
    }
    public boolean checkColor(RBTreeNode root){
        if (root == null) return true;

        if (root.color == Color.RED){
            RBTreeNode parent = root.parent;
            if (parent.color == Color.RED){
                System.out.println("连续两个红色节点!");
                return false;
            }
        }
        checkColor(root.left);
        checkColor(root.right);

        return true;
    }

    /**
     * @param pathBlackNum 走到当前路径的黑色节点个数
     * @param blackNum 每条路径该有的黑色节点个数
     */
    public boolean checkBlackNumber(RBTreeNode root, int pathBlackNum, int blackNum){
        if (root == null) return true;
        if (root.color == Color.BLACK){
            pathBlackNum++;
        }
        if (root.left == null && root.right == null) return pathBlackNum == blackNum;

        if (!checkBlackNumber(root.left, pathBlackNum, blackNum)){
            return false;
        }

        return checkBlackNumber(root.right, pathBlackNum, blackNum);
    }

    /**
     * 中序遍历
     */
    public void inorder(RBTreeNode root){
        if (root == null) return;

        inorder(root.left);
        System.out.print(root.val + "  ");
        inorder(root.right);
    }
}
